// Copyright 2022 jeff.li. and/or its affiliates.


#include <tbir/runtime/container.h>
#include <tbir/runtime/container/ndarray_helper.h>
#include <tbir/runtime/container_private.h>
#include <tbir/runtime/device_api.h>
#include <tbir/runtime/generic/generic_constructor_funcs.h>
#include <tbir/runtime/generic/generic_funcs.h>
#include <tbir/runtime/registry.h>

namespace tbir::runtime {

    TBIR_REGISTER_GLOBAL("runtime.NDArray").set_body([](PyArgs args) -> RTValue {
        RTValue data = args[0].As<RTValue>();
        List shape = args[1].As<List>();
        Unicode dtype_str = args[2].As<Unicode>();
        Unicode device_str = args[3].As<Unicode>();
        return Kernel_NDArray::make(data, shape, dtype_str, device_str);
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayToList").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.ToList();
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayIsContiguous").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.IsContiguous();
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayContiguous").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.Contiguous();
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayReshape").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.Reshape(args[1]);
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArraySqueeze").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.Squeeze(args[1]);
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayUnsqueeze").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.Unsqueeze(args[1].As<int64_t>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayStride").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        const int64_t *strides = data.GetStridesPtr();
        return List(strides, strides + data.GetDim());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayGetItem").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.get_item(args[1].As<RTValue>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArraySetItem").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        data.set_item(args[1].As<RTValue>(), args[2].As<RTValue>());
        return None;
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayGetSlice").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.get_slice(args[1].As<int64_t>(), args[2].As<int64_t>(), args[3].As<int64_t>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArraySetSlice").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        data.set_slice(args[1].As<int64_t>(), args[2].As<int64_t>(), args[3].As<RTValue>());
        return None;
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayTranspose").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.transpose(args[1].As<RTValue>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayAsType").set_body([](PyArgs args) -> RTValue {
        NDArray data = args[0].As<NDArray>();
        return data.as_type(args[1].As<Unicode>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayAdd").set_body([](PyArgs args) -> RTValue {
        return kernel_nd_module_add(args[0].As<RTView>(), args[1].As<RTView>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArraySub").set_body([](PyArgs args) -> RTValue {
        return kernel_nd_module_sub(args[0].As<RTView>(), args[1].As<RTView>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayDiv").set_body([](PyArgs args) -> RTValue {
        return kernel_nd_module_div(args[0].As<RTView>(), args[1].As<RTView>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayMul").set_body([](PyArgs args) -> RTValue {
        return kernel_nd_module_mul(args[0].As<RTView>(), args[1].As<RTView>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayRand").set_body([](PyArgs args) -> RTValue {
        return kernel_nd_module_rand(args[0].As<RTView>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayConcatenate").set_body([](PyArgs args) -> RTValue {
        return NDArrayOperate::Concatenate(args[0].As<RTValue>(), args[1].As<int64_t>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayStack").set_body([](PyArgs args) -> RTValue {
        return NDArrayOperate::Stack(args[0].As<RTValue>(), args[1].As<int64_t>());
    });

    TBIR_REGISTER_GLOBAL("runtime.NDArrayCopyToBytes").set_body([](PyArgs args) -> RTValue {
        void *to = reinterpret_cast<void *>(args[0].As<int64_t>());
        auto view = args[1].AsObjectView<NDArray>();
        auto &nd = view.data();
        TbirStreamHandle stream = reinterpret_cast<TbirStreamHandle>(args[2].As<int64_t>());
        MXCHECK(nd.IsContiguous()) << "NDArrayCopyToBytes: only support contiguous NDArray";
        // DeviceAPI::Get(nd->device)->DefaultComputeStreamSync(nd->device);
        DeviceAPI::Get(nd->device)
                ->CopyDataFromTo(nd->data,
                                 nd->byte_offset,
                                 to,
                                 0,
                                 nd.DataSize(),
                                 nd->device,
                                 nd->device,
                                 nd->dtype,
                                 stream);
        return None;
    });

}  // namespace tbir::runtime
